#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
公司信息OCR处理工具 - 基于通义千问多模态API

这个脚本专门用于处理公司信息图像的OCR识别，并以JSON格式输出。
支持命令行处理模式和API服务模式。
"""

import os
import sys
import json
import argparse
from pathlib import Path
from rich.console import Console
from rich.markdown import Markdown
import base64
import time

# FastAPI相关依赖
from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, HTTPException
from fastapi.responses import JSONResponse
import uvicorn
from typing import List, Optional
import tempfile
import uuid
import shutil
from pydantic import BaseModel

# 导入通义千问多模态工具
try:
    from qwen_multimodal import query_model, encode_image_to_base64
except ImportError:
    script_dir = os.path.dirname(os.path.abspath(__file__))
    sys.path.append(script_dir)
    from qwen_multimodal import query_model, encode_image_to_base64

console = Console()

# 添加debug模式标志
debug_mode = False

# 创建路由器
router_company = APIRouter(
    prefix="/company",
    tags=["公司信息OCR"],
    responses={404: {"description": "Not found"}},
)

# 构建提示词
def build_prompt():
    """构建用于公司信息识别的提示词"""
    prompt = """提取图片中公司信息为JSON格式。严格要求:
1. 识别图片中所有关于公司的信息，提取以下字段：
   - 统一社会信用代码：公司的统一社会信用代码，通常为18位字符
   - 公司名称：公司的完整名称
   - 电话：公司的联系电话
   - 企查查行业：公司所属行业分类
   - 企业规模：如大型、中型、小型等
   - 注册资本：公司的注册资本金额
   - 员工人数（2024年）：公司当前的员工数量
   - 成立日期：公司的成立日期，格式如YYYY-MM-DD
   - 地址：公司的详细地址
   - 简介：【重要】公司的详细介绍，必须完整提取全部内容，不要省略、简化或截断任何部分，不允许出现"收起"、"展开"等控制标记或省略号
   - 企查分：企查查平台上的企业评分
   - 科创分：科创相关评分
2. 对于没有找到的信息，使用"-"填充，不要编造或猜测任何信息
3. 不要添加图片中没有明确显示的任何内容
4. 返回JSON结构必须使用统一社会信用代码作为主键
5. 所有字段必须使用中文字段名
6. 对于简介内容：
   - 必须完整提取全部文本，不要截断，不能省略内容
   - 不要添加省略号、"收起"、"查看更多"等控制标记
   - 如果简介过长，仍然需要完整提取，不允许简化或截断
7. JSON结构必须符合以下格式:
{
  "统一社会信用代码值": {
    "统一社会信用代码": "统一社会信用代码值",
    "公司名称": "公司名称值",
    "电话": "电话值",
    "企查查行业": "行业值",
    "企业规模": "规模值",
    "注册资本": "资本值",
    "员工人数（2024年）": "人数值",
    "成立日期": "日期值",
    "地址": "地址值",
    "简介": "完整的公司简介内容，不省略不截断",
    "企查分": "分数值",
    "科创分": "分数值"
  }
}"""
    
    return prompt

# 调用API
def call_api(prompt, image_path_or_base64):
    """调用模型API进行图像识别，增强错误处理和备选方案"""
    try:
        # 检查图像路径是否有效
        if isinstance(image_path_or_base64, str) and not image_path_or_base64.startswith("data:image"):
            # 检查文件是否超过大小限制
            try:
                file_size = os.path.getsize(image_path_or_base64) / (1024 * 1024)  # 转换为MB
                if file_size > 10:  # 如果图像超过10MB
                    console.print(f"[yellow]警告: 图像大小为 {file_size:.2f}MB，超过10MB可能导致API失败[/yellow]")
                    # 尝试压缩图像
                    try:
                        from PIL import Image
                        compressed_path = compress_image(image_path_or_base64)
                        if compressed_path:
                            console.print(f"[green]已将图像压缩并保存至: {compressed_path}[/green]")
                            image_path_or_base64 = compressed_path
                    except ImportError:
                        console.print("[yellow]缺少PIL库，无法压缩图像[/yellow]")
            except Exception as e:
                console.print(f"[yellow]检查图像大小时出错: {e}[/yellow]")

            # 尝试使用直接API调用方式
            try:
                try:
                    from qwen_multimodal import direct_api_call
                    console.print("[yellow]尝试使用直接API调用...[/yellow]")
                    response = direct_api_call(prompt, image_path_or_base64)
                    if response:
                        return response
                except Exception as e:
                    console.print(f"[yellow]直接API调用失败: {e}，尝试使用query_model...[/yellow]")
                
                # 尝试使用query_model
                try:
                    response = query_model(prompt, image_path_or_base64)
                    if response:
                        return response
                except Exception as e:
                    console.print(f"[yellow]query_model调用失败: {e}[/yellow]")
                
                # 如果两种方法都失败，尝试使用base64编码
                console.print("[yellow]常规API调用方法均失败，尝试使用base64编码...[/yellow]")
                try:
                    base64_image = encode_image_to_base64(image_path_or_base64)
                    from qwen_multimodal import direct_api_call
                    response = direct_api_call(prompt, base64_image)
                    return response
                except Exception as e:
                    console.print(f"[yellow]base64编码方法也失败: {e}[/yellow]")
                    raise
            except Exception as e:
                console.print(f"[red]所有API调用方法均失败: {e}[/red]")
                raise
        else:
            # 如果已经是base64字符串，直接尝试调用
            try:
                from qwen_multimodal import direct_api_call
                response = direct_api_call(prompt, image_path_or_base64)
                return response
            except Exception as e:
                console.print(f"[yellow]直接API调用失败: {e}[/yellow]")
                try:
                    response = query_model(prompt, image_path_or_base64)
                    return response
                except Exception as e:
                    console.print(f"[yellow]query_model调用也失败: {e}[/yellow]")
                    raise
        
        # 如果代码执行到这里，表示所有调用方法都失败了
        raise Exception("所有API调用方法都失败")
    except Exception as e:
        print(f"API调用失败: {e}")
        if debug_mode:
            import traceback
            traceback.print_exc()
        return None

def compress_image(image_path, quality=85, max_size=8.0):
    """压缩图像以减小文件大小
    
    Args:
        image_path: 原始图像路径
        quality: 压缩质量，1-100
        max_size: 压缩后的最大大小(MB)
        
    Returns:
        str: 压缩后的图像路径，如果压缩失败则返回None
    """
    try:
        from PIL import Image
        import os
        
        # 生成压缩图像的文件名
        path, ext = os.path.splitext(image_path)
        compressed_path = f"{path}_compressed{ext}"
        
        # 打开图像
        img = Image.open(image_path)
        
        # 保存压缩版本
        img.save(compressed_path, quality=quality, optimize=True)
        
        # 检查压缩后的大小
        compressed_size = os.path.getsize(compressed_path) / (1024 * 1024)  # 转换为MB
        
        # 如果仍然太大且图像较大，可以调整尺寸
        if compressed_size > max_size and (img.width > 1000 or img.height > 1000):
            # 计算缩放因子，尝试达到目标大小
            scale_factor = 0.8  # 初始缩放因子
            
            while compressed_size > max_size and scale_factor > 0.3:  # 限制最小缩放到原来的30%
                # 计算新尺寸
                new_width = int(img.width * scale_factor)
                new_height = int(img.height * scale_factor)
                
                # 调整图像大小
                resized_img = img.resize((new_width, new_height), Image.LANCZOS)
                
                # 保存调整大小后的图像
                resized_img.save(compressed_path, quality=quality, optimize=True)
                
                # 检查新文件大小
                compressed_size = os.path.getsize(compressed_path) / (1024 * 1024)
                
                # 减小缩放因子，为下一次迭代准备
                scale_factor *= 0.9
        
        return compressed_path
    except Exception as e:
        console.print(f"[red]压缩图像失败: {e}[/red]")
        return None

# 从响应中提取JSON
def extract_json_from_response(response):
    """从API响应中提取JSON内容"""
    if not response:
        return None
    
    try:
        # 先尝试简单的提取方式
        json_start = response.find("{")
        json_end = response.rfind("}")
        
        if json_start != -1 and json_end != -1:
            json_text = response[json_start:json_end+1]
            try:
                result = json.loads(json_text)
                return result
            except json.JSONDecodeError as e:
                print(f"简单提取后解析JSON失败: {e}")
                # 尝试修复常见的JSON结构问题
                json_text = fix_json_structure(json_text)
                try:
                    result = json.loads(json_text)
                    return result
                except:
                    pass  # 如果修复后仍无法解析，继续尝试下一种方法
        
        # 如果简单提取失败，尝试使用正则表达式提取完整JSON
        import re
        json_pattern = r'\{(?:[^{}]|(?:\{[^{}]*\}))*\}'
        matches = re.findall(json_pattern, response)
        
        if matches:
            # 尝试解析找到的最长匹配项
            matches.sort(key=len, reverse=True)
            for match in matches:
                try:
                    result = json.loads(match)
                    return result
                except:
                    continue
        
        # 所有提取方法都失败
        print("无法从响应中提取有效的JSON")
        return None
    except Exception as e:
        print(f"提取JSON过程中出错: {e}")
        if debug_mode:
            import traceback
            traceback.print_exc()
        return None

def fix_json_structure(json_text):
    """尝试修复常见的JSON结构问题"""
    import re
    
    # 确保JSON对象完整闭合
    if not json_text.strip().endswith("}"):
        json_text = json_text.rstrip() + "}"
    
    # 修复缺少引号的键
    keys_pattern = r'(\s*)([a-zA-Z\u4e00-\u9fa5]+)(\s*):(\s*)'
    json_text = re.sub(keys_pattern, r'\1"\2"\3:\4', json_text)
    
    # 修复可能缺少的逗号
    comma_pattern = r'([\"\d])\s*\n\s*\"'
    json_text = re.sub(comma_pattern, r'\1,\n"', json_text)
    
    return json_text

def clean_company_info(result):
    """清理和验证公司信息结果"""
    if not isinstance(result, dict):
        return result
    
    cleaned_result = {}
    
    # 遍历结果中的每个公司
    for code, company_info in result.items():
        if not isinstance(company_info, dict):
            continue
        
        # 确保所有必要字段都存在
        required_fields = [
            "统一社会信用代码", "公司名称", "电话", "企查查行业", "企业规模",
            "注册资本", "员工人数（2024年）", "成立日期", "地址", "简介",
            "企查分", "科创分"
        ]
        
        clean_company = {}
        for field in required_fields:
            # 如果字段不存在或为空，设置为"-"
            if field not in company_info or not company_info[field] or company_info[field].strip() == "":
                clean_company[field] = "-"
            else:
                # 清理字段值中的多余空格
                value = company_info[field].strip()
                
                # 特殊处理简介字段，确保不省略内容且移除末尾的"收起"标记
                if field == "简介":
                    # 移除简介末尾的"收起"或"展开"等控制标记
                    value = value.replace("收起", "").replace("展开", "")
                    value = value.replace("查看更多", "").replace("查看全部", "")
                    # 确保结尾没有不必要的符号
                    value = value.rstrip("。，,. ") + "。" if value and value != "-" else value
                
                clean_company[field] = value
        
        # 确保统一社会信用代码作为主键一致
        clean_company["统一社会信用代码"] = code
        
        # 添加到清理后的结果
        cleaned_result[code] = clean_company
    
    return cleaned_result

def process_single_image(image_path, retries=3):
    """处理单张公司信息图片并返回结构化信息
    
    Args:
        image_path: 图片的路径
        retries: 重试次数，默认为3次
        
    Returns:
        dict: 结构化的公司信息
    """
    # 检查图片路径是否存在
    if not os.path.exists(image_path):
        print(f"错误: 无法找到图片文件 {image_path}")
        return None
    
    # 构建提示词
    prompt = build_prompt()
    
    console.print(f"[bold]正在处理公司信息图像: {image_path}[/bold]")
    console.print("[bold blue]正在识别公司信息...[/bold blue]")
    
    last_error = None
    best_result = None
    
    for attempt in range(retries):
        try:
            # 针对不同重试次数使用不同的策略
            if attempt > 0:
                console.print(f"[yellow]重试第 {attempt+1} 次...[/yellow]")
                
                # 根据重试次数调整提示词
                if attempt == 1:
                    retry_prompt = f"{prompt}\n\n这是重试。请特别仔细识别所有公司相关信息，特别是统一社会信用代码和公司名称。"
                else:
                    retry_prompt = "请从图片中提取公司关键信息，尤其关注统一社会信用代码、公司名称、电话、地址等基本信息。对于无法确定的内容使用'-'表示，不要猜测。请以JSON格式返回结果。"
                
                # 调整API调用方法
                if attempt == 1:
                    # 第二次尝试，直接使用原始图像
                    response = call_api(retry_prompt, image_path)
                elif attempt == 2:
                    # 第三次尝试，尝试压缩图像
                    try:
                        compressed_path = compress_image(image_path)
                        if compressed_path:
                            console.print(f"[yellow]使用压缩图像重试: {compressed_path}[/yellow]")
                            response = call_api(retry_prompt, compressed_path)
                        else:
                            response = call_api(retry_prompt, image_path)
                    except Exception:
                        response = call_api(retry_prompt, image_path)
                else:
                    # 更多重试，尝试其他方法
                    try:
                        # 尝试使用base64编码
                        base64_image = encode_image_to_base64(image_path)
                        response = call_api(retry_prompt, base64_image)
                    except Exception:
                        response = call_api(retry_prompt, image_path)
            else:
                # 第一次尝试，使用原始提示词
                first_prompt = f"{prompt}\n\n请特别注意识别统一社会信用代码和公司名称，这些是最关键的字段。对于简介字段，必须完整输出全部内容，不要省略任何部分，不要截断。对于没有明确看到的信息，请使用'-'填充，不要编造。"
                response = call_api(first_prompt, image_path)
            
            # 处理API响应
            if response:
                result = extract_json_from_response(response)
                
                # 如果提取失败但还有重试机会
                if not result and attempt < retries - 1:
                    console.print("[bold red]无法提取JSON结果，将重试...[/bold red]")
                    continue
                
                # 如果有结果，进行处理
                if result:
                    # 特殊处理简介字段，确保不省略内容
                    for code, company in result.items():
                        if isinstance(company, dict) and "简介" in company:
                            # 移除简介末尾的"收起"或"展开"等控制标记
                            intro = company["简介"]
                            if intro and intro != "-":
                                # 移除各种可能导致内容被截断的标记
                                intro = intro.replace("收起", "").replace("展开", "")
                                intro = intro.replace("查看更多", "").replace("查看全部", "")
                                intro = intro.replace("点击查看更多", "").replace("点击展开全部", "")
                                intro = intro.replace("展开全部", "").replace("收起全部", "")
                                intro = intro.replace("展开详情", "").replace("收起详情", "")
                                intro = intro.replace("...更多", "").replace("... ...", "")
                                intro = intro.replace("……", "").replace("...", "")
                                
                                # 检查简介是否以省略号结尾，若是则提示需特别注意
                                if intro.rstrip().endswith(("...", "…", "....")) or "..." in intro:
                                    console.print("[bold yellow]警告: 简介可能被截断，请检查原始图片确认完整内容[/bold yellow]")
                                
                                # 确保结尾没有不必要的符号
                                intro = intro.rstrip("。，,. ") + "。" if intro and intro != "-" else intro
                                company["简介"] = intro
                                
                    # 清理和验证公司信息
                    result = clean_company_info(result)
                    
                    # 更新最佳结果
                    best_result = result
                
                # 如果是最后一次尝试或已有很好的结果，就终止
                if attempt == retries - 1 or best_result:
                    if best_result:
                        console.print("[bold green]识别完成！[/bold green]")
                        return best_result
                    
            elif attempt < retries - 1:
                # 如果没有响应但还有重试机会
                console.print("[bold yellow]未获取有效响应，将重试...[/bold yellow]")
                continue
                
        except Exception as e:
            last_error = e
            console.print(f"[bold red]处理图片时出错: {str(e)}[/bold red]")
            if debug_mode:
                import traceback
                traceback.print_exc()
            
            if attempt < retries - 1:
                # 还有重试机会
                continue
    
    # 所有尝试都失败且没有任何有效结果
    if last_error:
        console.print(f"[bold red]所有重试都失败，最后错误: {str(last_error)}[/bold red]")
    
    return None

def process_directory(directory_path, output_path=None, retries=3):
    """处理目录中的所有图片
    
    Args:
        directory_path: 包含公司信息图片的目录路径
        output_path: 保存结果的JSON文件路径
        retries: 重试次数
        
    Returns:
        dict: 包含所有公司信息的字典
    """
    # 检查目录是否存在
    if not os.path.isdir(directory_path):
        console.print(f"[bold red]错误: 目录不存在 {directory_path}[/bold red]")
        return None
    
    # 获取所有图片文件
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff']
    image_files = []
    
    for root, _, files in os.walk(directory_path):
        for file in files:
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(root, file))
    
    if not image_files:
        console.print(f"[bold yellow]警告: 目录中没有找到图片文件 {directory_path}[/bold yellow]")
        return {}
    
    console.print(f"[bold]找到 {len(image_files)} 个图片文件[/bold]")
    
    # 处理所有图片
    all_results = {}
    
    for i, image_file in enumerate(image_files):
        console.print(f"[bold]处理图片 {i+1}/{len(image_files)}: {image_file}[/bold]")
        
        result = process_single_image(image_file, retries)
        
        if result:
            # 合并结果
            all_results.update(result)
            
            # 对所有结果中的简介字段进行处理，确保不省略内容
            for code, company in all_results.items():
                if isinstance(company, dict) and "简介" in company:
                    # 移除简介末尾的"收起"或"展开"等控制标记
                    intro = company["简介"]
                    if intro and intro != "-":
                        # 移除各种可能导致内容被截断的标记
                        intro = intro.replace("收起", "").replace("展开", "")
                        intro = intro.replace("查看更多", "").replace("查看全部", "")
                        intro = intro.replace("点击查看更多", "").replace("点击展开全部", "")
                        intro = intro.replace("展开全部", "").replace("收起全部", "")
                        intro = intro.replace("展开详情", "").replace("收起详情", "")
                        intro = intro.replace("...更多", "").replace("... ...", "")
                        intro = intro.replace("……", "").replace("...", "")
                        
                        # 检查简介是否以省略号结尾，若是则提示需特别注意
                        if intro.rstrip().endswith(("...", "…", "....")) or "..." in intro:
                            console.print("[bold yellow]警告: 简介可能被截断，请检查原始图片确认完整内容[/bold yellow]")
                        
                        # 确保结尾没有不必要的符号
                        intro = intro.rstrip("。，,. ") + "。" if intro and intro != "-" else intro
                        company["简介"] = intro
            
            # 保存中间结果（每处理完一张图片都保存一次）
            if output_path:
                try:
                    with open(output_path, 'w', encoding='utf-8') as f:
                        json.dump(all_results, f, ensure_ascii=False, indent=2)
                    console.print(f"[green]中间结果已保存到: {output_path}[/green]")
                except Exception as e:
                    console.print(f"[red]保存中间结果时出错: {e}[/red]")
    
    console.print(f"[bold green]所有图片处理完成，共提取到 {len(all_results)} 家公司的信息[/bold green]")
    
    # 保存最终结果
    if output_path and all_results:
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(all_results, f, ensure_ascii=False, indent=2)
            console.print(f"[bold green]最终结果已保存到: {output_path}[/bold green]")
        except Exception as e:
            console.print(f"[bold red]保存最终结果时出错: {str(e)}[/bold red]")
    
    return all_results

# 删除临时文件的后台任务
def remove_temp_file(file_path: str):
    """删除临时文件的后台任务"""
    try:
        if os.path.exists(file_path):
            if os.path.isdir(file_path):
                shutil.rmtree(file_path)
            else:
                os.remove(file_path)
    except Exception as e:
        console.print(f"[yellow]警告: 无法删除临时文件 {file_path}: {e}[/yellow]")

# API路由

# 健康检查端点
@router_company.get("/health")
async def health_check():
    """健康检查接口，用于监控服务是否正常运行"""
    return {"status": "healthy", "service": "company-ocr-api", "timestamp": time.time()}

# 添加支持JSON请求体的模型类
class FilePathRequest(BaseModel):
    file_path: str
    retries: Optional[int] = 3
    save_path: Optional[str] = None

@router_company.post("/ocr_json")
async def process_image_json_api(request: FilePathRequest):
    """通过JSON处理单张图片或文件夹
    
    Args:
        file_path: 图片文件路径或文件夹路径
        retries: 识别失败时的重试次数，默认为3次
        save_path: 保存JSON结果的路径，默认为file_path所在目录下的ocr_result.json
    
    Returns:
        提取的公司信息，JSON格式
    """
    path = request.file_path
    
    # 检查路径是否存在
    if not os.path.exists(path):
        return JSONResponse(
            status_code=400,
            content={"error": f"指定的路径不存在: {path}"}
        )
    
    try:
        # 设置默认保存路径
        save_path = request.save_path
        if not save_path:
            # 如果是文件夹，保存在文件夹内
            if os.path.isdir(path):
                save_path = os.path.join(path, "ocr_result.json")
            else:
                # 如果是文件，保存在文件所在目录
                save_dir = os.path.dirname(path)
                save_path = os.path.join(save_dir, "ocr_result.json")
        
        # 判断是文件还是文件夹
        if os.path.isdir(path):
            # 如果是文件夹，处理文件夹中的所有图片
            result = process_directory(path, output_path=save_path, retries=request.retries)
        else:
            # 如果是文件，处理单个文件
            result = process_single_image(path, retries=request.retries)
            # 保存结果
            if result and save_path:
                try:
                    # 确保保存目录存在
                    save_dir = os.path.dirname(save_path)
                    os.makedirs(save_dir, exist_ok=True)
                    # 保存为JSON文件
                    with open(save_path, 'w', encoding='utf-8') as f:
                        json.dump(result, f, ensure_ascii=False, indent=2)
                    console.print(f"[green]结果已保存到: {save_path}[/green]")
                except Exception as e:
                    console.print(f"[red]保存结果时出错: {e}[/red]")
        
        if not result:
            return JSONResponse(
                status_code=500,
                content={"error": "图像处理失败，无法提取公司信息"}
            )
        
        # 在返回结果前，对所有公司的简介字段进行最终检查处理
        if result and isinstance(result, dict):
            for code, company in result.items():
                if isinstance(company, dict) and "简介" in company and company["简介"] and company["简介"] != "-":
                    intro = company["简介"]
                    # 再次清理所有可能的截断标记
                    intro = intro.replace("收起", "").replace("展开", "")
                    intro = intro.replace("查看更多", "").replace("查看全部", "")
                    intro = intro.replace("点击查看更多", "").replace("点击展开全部", "")
                    intro = intro.replace("展开全部", "").replace("收起全部", "")
                    intro = intro.replace("展开详情", "").replace("收起详情", "")
                    intro = intro.replace("...更多", "").replace("... ...", "")
                    intro = intro.replace("……", "").replace("...", "")
                    # 确保结尾有标点
                    intro = intro.rstrip("。，,. ") + "。" if intro else intro
                    company["简介"] = intro
        
        # 在返回结果中添加保存路径信息
        result_with_path = {
            "data": result,
            "save_path": save_path if os.path.exists(save_path) else None
        }
        
        return result_with_path
        
    except Exception as e:
        console.print(f"[bold red]处理请求时出错: {str(e)}[/bold red]")
        if debug_mode:
            import traceback
            traceback.print_exc()
        return JSONResponse(
            status_code=500,
            content={"error": f"处理请求时出错: {str(e)}"}
        )

# 提供命令行入口
def main():
    parser = argparse.ArgumentParser(description='公司信息OCR处理工具')
    parser.add_argument('--folder', type=str, help='输入图片目录路径')
    parser.add_argument('--output', '-o', type=str, help='输出JSON文件路径')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    parser.add_argument('--retries', '-r', type=int, default=3, help='识别失败时的重试次数')
    
    args = parser.parse_args()
    
    # 设置调试模式
    global debug_mode
    if args.debug:
        debug_mode = True
        print("调试模式已启用")
    
    # 命令行处理模式
    if not args.folder:
        parser.error("命令行处理模式需要提供 --folder 参数指定输入图片目录或文件路径")
    
    # 确定默认输出路径
    output_path = args.output
    if not output_path:
        if os.path.isdir(args.folder):
            output_path = os.path.join(args.folder, "company_info_results.json")
        else:
            output_dir = os.path.dirname(args.folder)
            if not output_dir:
                output_dir = "."
            output_path = os.path.join(output_dir, "company_info_results.json")
    
    # 处理输入
    if os.path.isdir(args.folder):
        # 如果输入是目录，处理目录中的所有图片
        results = process_directory(args.folder, output_path, args.retries)
    else:
        # 如果输入是单个文件，只处理该文件
        result = process_single_image(args.folder, args.retries)
        if result:
            # 保存结果
            if output_path:
                try:
                    with open(output_path, 'w', encoding='utf-8') as f:
                        json.dump(result, f, ensure_ascii=False, indent=2)
                    console.print(f"[bold green]结果已保存到: {output_path}[/bold green]")
                except Exception as e:
                    console.print(f"[bold red]保存结果时出错: {str(e)}[/bold red]")
            results = result
    
    return results

if __name__ == "__main__":
    main()